#include "kernel_operator.h"

constexpr uint32_t FLOAT_DATA_TYPE_SIZE = 4;                        
constexpr uint32_t FLOAT64_DATA_TYPE_SIZE = 8;                         
constexpr uint32_t BLOCK_SIZE = 32;                                         
constexpr uint32_t TILE_ELEM_NUM = BLOCK_SIZE / FLOAT_DATA_TYPE_SIZE;             // float的Tile元素数（8）
constexpr uint32_t TILE_ELEM_NUM_FLOAT64 = BLOCK_SIZE / FLOAT64_DATA_TYPE_SIZE;  
constexpr uint32_t BUFFER_NUM = 2;                                                // 双缓冲

template <typename TStart, typename TEnd> 
class KernelLinSpaceD {
public:
    __aicore__ inline KernelLinSpaceD() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, 
                                uint32_t totalLength, uint32_t formerLength,
                                uint32_t formerTileNum, uint32_t formerLastTileLength,
                                uint32_t formerNum, uint32_t tailLength,
                                uint32_t tailTileNum, uint32_t tailLastTileLength) 
    {
        startGm.SetGlobalBuffer((__gm__ TStart*)x, 1);
        endGm.SetGlobalBuffer((__gm__ TEnd*)y, 1);
        this->start = startGm.GetValue(0);
        this->end = endGm.GetValue(0);
        this->origSize = totalLength;

        uint32_t coreId = AscendC::GetBlockIdx();
        uint32_t currentBlockLength = 0; 

        if (coreId < formerNum) 
        {
            this->blockLength = formerLength;
            this->tileNum = formerTileNum;
            this->lastTileLength = formerLastTileLength;
            this->blockOffset = formerLength * coreId;
            currentBlockLength = formerLength;
        } 
        else 
        {
            this->blockLength = tailLength;
            this->tileNum = tailTileNum;
            this->lastTileLength = tailLastTileLength;
            this->blockOffset = formerLength * formerNum + tailLength * (coreId - formerNum);
            currentBlockLength = tailLength;
        }    
        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(z) + this->blockOffset, currentBlockLength);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, TILE_ELEM_NUM * sizeof(float));
        pipe.InitBuffer(tmpHalf, sizeof(half));
        pipe.InitBuffer(tmpFloat, sizeof(float));     
        pipe.InitBuffer(tmpStart, sizeof(TStart));
        pipe.InitBuffer(tmpEnd, sizeof(TEnd));     
    }
    __aicore__ inline void Process()
    {
        if constexpr (AscendC::Std::is_same<TStart, bfloat16_t>::value) {
            this->startF = AscendC::ToFloat(this->start);
        }
        else if constexpr (AscendC::Std::is_same<TStart, uint8_t>::value) {
            AscendC::LocalTensor<TStart> uint8Local = tmpStart.Get<TStart>();
            AscendC::LocalTensor<half> halfLocal = tmpHalf.Get<half>();
            AscendC::LocalTensor<float> floatLocal = tmpFloat.Get<float>();
            uint8Local.SetValue(0, this->start);
            Cast(halfLocal, uint8Local, AscendC::RoundMode::CAST_NONE, 1);
            Cast(floatLocal, halfLocal, AscendC::RoundMode::CAST_NONE, 1);
            this->startF = floatLocal.GetValue(0);
        }          
        else if constexpr (AscendC::Std::is_same<TStart, float>::value) {
            this->startF = this->start;
        }   
        else {
            this->startF = static_cast<float>(this->start);
        }   

        if constexpr (AscendC::Std::is_same<TEnd, bfloat16_t>::value) {
            this->endF = AscendC::ToFloat(this->end);
        }
        else if constexpr (AscendC::Std::is_same<TEnd, uint8_t>::value) {
            AscendC::LocalTensor<half> halfLocal = tmpHalf.Get<half>();
            AscendC::LocalTensor<float> floatLocal = tmpFloat.Get<float>();
            AscendC::LocalTensor<TEnd> uint8Local = tmpEnd.Get<TEnd>();
            uint8Local.SetValue(0, this->end);
            Cast(halfLocal, uint8Local, AscendC::RoundMode::CAST_NONE, 1);
            Cast(floatLocal, halfLocal, AscendC::RoundMode::CAST_NONE, 1);
            this->endF = floatLocal.GetValue(0);
        }      
        else if constexpr (AscendC::Std::is_same<TStart, float>::value) {
            this->endF = this->end;
        }      
        else {
            this->endF = static_cast<float>(this->end);
        }  

        if (this->origSize == 1) 
        {
            uint32_t coreId = AscendC::GetBlockIdx();
            if(coreId >= 1) return;
            AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
            AscendC::Duplicate<float>(tileOutput, startF, TILE_ELEM_NUM);
            outQueueZ.EnQue(tileOutput);
            CopyOut(0);
        }
        else{
            this->step = (endF - startF) / (origSize - 1);
            for (int32_t tileIdx = 0; tileIdx < tileNum; tileIdx++) 
            {
                Compute(tileIdx);
                CopyOut(tileIdx);
            }
        }
    }

private:
    __aicore__ inline void Compute(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
        int32_t tileElemNum = TILE_ELEM_NUM;  // 8
        int32_t tileGlobalStart = blockOffset + tileIdx * tileElemNum;
        for (int32_t i = tileElemNum - 1; i >= 0; i--) {
            AscendC::Duplicate<float>(tileOutput, startF + step * static_cast<float>(tileGlobalStart + i), i + 1);
        }
        outQueueZ.EnQue(tileOutput);
    }
    __aicore__ inline void CopyOut(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        int32_t copyElemNum;                         
        if (tileIdx == tileNum - 1) {
            copyElemNum = lastTileLength;
        }
        else {
            copyElemNum = TILE_ELEM_NUM;
        }

        int32_t copyOffset = tileIdx * TILE_ELEM_NUM;
        if (copyOffset + copyElemNum > blockLength) {
            copyElemNum = blockLength - copyOffset;
        }
        AscendC::DataCopy(zGm[copyOffset], zLocal, copyElemNum);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tmpHalf, tmpFloat;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tmpStart, tmpEnd;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<float> zGm;  
    AscendC::GlobalTensor<TStart> startGm;
    AscendC::GlobalTensor<TEnd> endGm;
    uint32_t blockLength;    // 当前核心处理的float总元素数
    uint32_t tileNum;        // 当前核心的Tile总数
    uint32_t lastTileLength; // 最后一个Tile的float元素数
    uint32_t blockOffset;    // 当前核心的全局float元素偏移
    uint32_t origSize;       // 原始有效元素数
    float step, startF, endF;
    TStart start;
    TEnd end;
};

extern "C" __global__ __aicore__ void lin_space_d_custom(GM_ADDR start, GM_ADDR end, GM_ADDR num, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);  
    if(TILING_KEY_IS(100)) { 
        // bfloat16, uint8, float16, float32, int8, int16, int32
        KernelLinSpaceD<DTYPE_START, DTYPE_END> op;
        op.Init(start, end, z,
                tiling_data.totalLength, tiling_data.formerLength,
                tiling_data.formerTileNum, tiling_data.formerLastTileLength,
                tiling_data.formerNum, tiling_data.tailLength,
                tiling_data.tailTileNum, tiling_data.tailLastTileLength);
        op.Process();
    }
/*
    else if(TILING_KEY_IS(101)) {
        // 预留float64接口
        KernelLinSpaceDouble<DTYPE_START, DTYPE_END> op;
        op.Init(start, end, z,
                tiling_data.totalLength, tiling_data.formerLength,
                tiling_data.formerTileNum, tiling_data.formerLastTileLength,
                tiling_data.formerNum, tiling_data.tailLength,
                tiling_data.tailTileNum, tiling_data.tailLastTileLength);
        op.Process();
    }
*/
}
